from langchain import OpenAI, SQLDatabase, SQLDatabaseChain
from langchain.prompts.prompt import PromptTemplate
import openai

_echarts_prompt = """You are a Echarts AND SQL expert. Given an input sql,You infer what type of ECharts chart to display, and return a syntactically correct Echarts config json.
Pay attention to use only the column names you can see in the SQL below. Be careful to not query for columns that do not exist.

Use the following format:

SQLQuery: {sql}
Type: type of ECharts
Echarts: Echarts Config

"""

ECHARTS_PROMPT = PromptTemplate(
    input_variables=["sql"],
    template=_echarts_prompt
)

db = SQLDatabase.from_uri(
    "mysql+pymysql://root:password@115.159.119.116:3306/openai_demo",
    include_tables=['data_wide_s_trade'],
    sample_rows_in_table_info=2
)

def message_to_sql(message):
    llm = OpenAI(temperature=0)
    db_chain = SQLDatabaseChain.from_llm(llm, db, verbose=True, use_query_checker=True, return_intermediate_steps=True, top_k=10)
    result = db_chain(message)
    return result["intermediate_steps"][1]   

def sql_to_charts(sql):
       p = ECHARTS_PROMPT.format(sql=sql)
       messages = [{"role": "system", "content": p}, {"role": "user", "content": sql}]
       response = openai.ChatCompletion.create(
            model="gpt-3.5-turbo",
            messages=messages, temperature=0.0, max_tokens=2048,
            n=1, presence_penalty=0, frequency_penalty=0
       )
       return response.choices[0].message.content

if __name__ == "__main__":
    sql=message_to_sql("今年每个项目签约了多少合同?")
    chat=sql_to_charts(sql)
    print(chat)